import numpy as np
import SimpleITK as sitk
from PIL import Image
import os
import shutil
import pandas as pd
from enum import Enum
import random
import torch

from mycode.pre_process import window_normalize
from mycode.section_synthesis import get_label_coord, draw_2Dsquare


secData_root = 'mydata_section'
landmark_root = os.path.join(secData_root, 'synthesis_landmark')
shortAxis_gt_root = os.path.join('mydata_gt', 'rotate', 'output_shortAxis')
shortAxis_pred_root = os.path.join('mydata_pred', 'rotate', 'output_shortAxis')
dataset_root = os.path.join(secData_root, 'myheart_hrnet_dataset')
image_root = os.path.join(dataset_root, 'images')
predlm_root = os.path.join(secData_root, 'pred_sparse_landmark')


SAX_PREFIX = 'case_'
SAX_NAME = 'imaging.nii.gz'
LM_PREFIX = 'synthesis_'
JPG_SUFFIX = '.jpg'
NIFTI_SUFFIX = '.nii.gz'
TRAIN_CSV = 'my_hrnet_train.csv'
TEST_CSV = 'my_hrnet_test.csv'
PRED_NAME = 'predictions.pth'
PRED_PREFIX = 'sparse_'


def generate_RGBImg(img_array: np.ndarray) -> Image:
    """
    generate the RGB Image of normalized 0-1 grey array, calculate the coordinate of landmarks in landmark_set.

    :param img_array: 2D 0-1 float32 grey array in CTA of order YX
    :return: Image: PIL Image to be saved
    """

    assert img_array.max() <= 1.0
    assert img_array.min() >= 0.0

    img_array *= 255
    img_array = np.array([img_array, img_array, img_array])
    img_array = img_array.transpose(1, 2, 0)
    img = Image.fromarray(img_array.astype(np.uint8)).convert('RGB')

    return img

class CSVAttri(Enum):
    ImageName = 'image_name'
    Scale = 'scale'
    CenterW = 'center_w'
    CenterH = 'center_h'
    UpperLm_X = 'original_0_x'
    UpperLm_Y = 'original_0_y'
    LowerLm_X = 'original_1_x'
    LowerLm_Y = 'original_1_y'


def select_img_from_cta(cta_array: np.ndarray, lm_array: np.ndarray, expect_num=40) -> (list, list, list):
    """
    select 2D slices from 3D CTA randomly

    :param cta_array: ZYX CTA ndarray
    :param lm_array: ZYX synthesis landmark ndarray
    :param expect_num: the number of slices to be selected if possible.
    :return: (list1, list2, list3), list1 contains Image type of the selected images,
            list2 contains two XY coordinates of the landmarks,
            list3 contains the z index of selected slices
    """
    lm_coords = get_label_coord(lm_array, {1.0, 2.0})
    upper_coords = lm_coords[1]
    lower_coords = lm_coords[2]
    start_z = max(upper_coords[0][0], lower_coords[0][0])
    end_z = min(upper_coords[-1][0], lower_coords[-1][0])
    expect_num = min(expect_num, end_z - start_z + 1)
    slices = random.sample(list(range(start_z, end_z + 1)), expect_num)
    slices.sort()

    cta_array = window_normalize(cta_array)
    upper_dict = dict({})
    lower_dict = dict({})
    # [[z,y,x]...] -> {z:[y,x],...}
    for coord in upper_coords:
        upper_dict[coord[0]] = coord[1:]
    for coord in lower_coords:
        lower_dict[coord[0]] = coord[1:]

    img_list = []
    coord_list = []
    for z in slices:
        img = generate_RGBImg(cta_array[z])
        coord = (upper_dict[z][::-1], lower_dict[z][::-1])
        img_list.append(img)
        coord_list.append(coord)
    return img_list, coord_list, slices


def dataset_creator(train_seq_list, test_seq_list, is_gt=False, is_append=False):
    """
    :param train_seq_list:
    :param test_seq_list:
    :param is_gt: whether the test dataset is created from groundtruth SAX
    :param is_append:
    :return:
    """

    shutil.rmtree(image_root)
    os.mkdir(image_root)

    train_items = []
    test_items = []
    total_seq_list = train_seq_list + test_seq_list


    for i in range(len(total_seq_list)):
        seq_number = total_seq_list[i]
        if is_gt or i < len(train_seq_list):
            shortAxis_root = shortAxis_gt_root
        else:
            shortAxis_root = shortAxis_pred_root
        cta_path = os.path.join(shortAxis_root, SAX_PREFIX + seq_number, SAX_NAME)
        lm_path = os.path.join(landmark_root, LM_PREFIX + seq_number + NIFTI_SUFFIX)
        cta_array = sitk.GetArrayFromImage(sitk.ReadImage(cta_path))
        lm_array = sitk.GetArrayFromImage(sitk.ReadImage(lm_path))

        Image_list, coord_list, selected_z = select_img_from_cta(cta_array, lm_array, 20)
        for j in range(len(Image_list)):
            # save the image
            image_name = seq_number + '_' + str(selected_z[j]) + JPG_SUFFIX
            Image_list[j].save(os.path.join(image_root, image_name))

            # create the csv item
            item = dict({})
            item[CSVAttri.ImageName.value] = image_name
            # the resolution of each Image is always 512X512
            item[CSVAttri.Scale.value] = 2.56
            item[CSVAttri.CenterW.value] = 255.5
            item[CSVAttri.CenterH.value] = 255.5
            item[CSVAttri.UpperLm_X.value] = coord_list[j][0][0]
            item[CSVAttri.UpperLm_Y.value] = coord_list[j][0][1]
            item[CSVAttri.LowerLm_X.value] = coord_list[j][1][0]
            item[CSVAttri.LowerLm_Y.value] = coord_list[j][1][1]

            items = train_items if i < len(train_seq_list) else test_items
            items.append(item)
        print('finished', seq_number, 'select and JPG create')

    df_train = pd.DataFrame(train_items, columns=[n.value for n in CSVAttri])
    df_test = pd.DataFrame(test_items, columns=[n.value for n in CSVAttri])
    df_train.to_csv(os.path.join(dataset_root, TRAIN_CSV), index=False)
    df_test.to_csv(os.path.join(dataset_root, TEST_CSV), index=False)


def parse_csv(df: pd.DataFrame) -> dict:
    """
    parse the csv file, return dict: k: seq_number, v:[z-index].
    the order is according with the csv row order.

    :param df:
    :return:
    """
    ret = dict({})
    for i, row in df.iterrows():
        name = row[CSVAttri.ImageName.value]
        seq_num, z_idx = name.split('.')[0].split('_')
        # coord0 = [row[CSVAttri.UpperLm_X.value], row[CSVAttri.UpperLm_Y.value]]
        # coord1 = [row[CSVAttri.LowerLm_X.value], row[CSVAttri.LowerLm_Y.value]]
        z_idx = int(z_idx)
        if seq_num not in ret:
            ret[seq_num] = []
        ret[seq_num].append(z_idx)
    return ret


# need to be tested
def pred_decoder(is_gt=False):
    """
    parse the model prediction output file .pth to sparse landmark into pred_sparse_landmark/
    is_gt: whether test dataset is created from groudtruth SAX
    :return:
    """
    pred = torch.load(os.path.join(dataset_root, PRED_NAME)).numpy()
    test_df = pd.read_csv(os.path.join(dataset_root, TEST_CSV))
    test_idx_dict = parse_csv(test_df)
    item_num = 0
    shortAxis_root = shortAxis_gt_root if is_gt else shortAxis_pred_root
    for seq_num, idx_list in test_idx_dict.items():
        sax_img = sitk.ReadImage(os.path.join(shortAxis_root, SAX_PREFIX + seq_num, SAX_NAME))
        pred_sparse = np.zeros(sax_img.GetSize())
        pred_sparse = pred_sparse.swapaxes(0, 2)
        for z_idx in idx_list:
            cur_upper_coord = pred[item_num][0].astype(np.uint32)
            cur_lower_coord = pred[item_num][1].astype(np.uint32)
            # print('seq num', seq_num)
            # print('z_idx', z_idx)
            # print('upper_coord zyx', (z_idx, cur_upper_coord[1], cur_upper_coord[0]))
            # print('lower_coord zyx', (z_idx, cur_lower_coord[1], cur_lower_coord[0]))
            draw_2Dsquare(pred_sparse, (z_idx, cur_upper_coord[1], cur_upper_coord[0]), label=1.0)  # the label value of upper point is 1.0
            draw_2Dsquare(pred_sparse, (z_idx, cur_lower_coord[1], cur_lower_coord[0]), label=2.0)
            item_num += 1
        pred_sparse_img = sitk.GetImageFromArray(pred_sparse)
        pred_sparse_img.CopyInformation(sax_img)
        sitk.WriteImage(pred_sparse_img, os.path.join(predlm_root, PRED_PREFIX + seq_num + NIFTI_SUFFIX))
        print('finished pred decoder:', seq_num)


def test():
    # test_img_path = r'/mytest/mycode/test_imaging_00184.nii.gz'
    # test_landmark_path = r'/mytest/mycode/test_synthesis_00184.nii.gz'
    # img_array = sitk.GetArrayFromImage(sitk.ReadImage(test_img_path))
    # img_array = window_normalize(img_array)
    # landmark_array = sitk.GetArrayFromImage(sitk.ReadImage(test_landmark_path))
    # coord_set = {1., 2.}
    #
    # img = generate_RGBImg(img_array[260])
    #
    # img.save('test_saveRGB.jpg')
    # img_array = np.array(Image.open('mycode/test_saveRGB.jpg').convert('RGB'), dtype=np.float32)
    # print(img_array.max())

    pred_decoder()


if __name__ == '__main__':

    train_list = [144, 146, 147, 149, 152, 154, 155, 158, 159, 160, 161, 162, 166, 167, 168, 170, 171, 172, 173, 176]
    train_list += list(range(200, 258))

    ori_test_list = list(range(177, 197))
    ori_test_list.remove(186)

    extra_test_list = list(range(260, 345))
    extra_test_list += list(range(350, 442))

    # train_seq_list = [str(i).zfill(5) for i in train_list]
    # test_seq_list = [str(i).zfill(5) for i in ori_test_list + extra_test_list]

    train_seq_list = ['00177', '00308']
    test_seq_list = ['00441']

    dataset_creator(train_seq_list, test_seq_list, is_gt=False)
    # pred_decoder(is_gt=False)

    pass

